package org.finesys.chat.core.base.embedding.store.milvus;

import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.*;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.QueryParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;
import org.finesys.chat.core.base.embedding.store.milvus.exception.RequestToMilvusFailedException;

import java.util.List;

public class CollectionOperationsExecutor {

    public CollectionOperationsExecutor() {

    }

    public static void flush(MilvusServiceClient client, String collectionName) {
        FlushParam flushParam = CollectionRequestBuilder.buildFlushRequest(collectionName);
        R<FlushResponse> response = client.flush(flushParam);
        checkResponseNotFailed(response);
    }

    static boolean hasCollection(MilvusServiceClient milvusClient, String collectionName) {
        HasCollectionParam request = CollectionRequestBuilder.buildHasCollectionRequest(collectionName);
        R<Boolean> response = milvusClient.hasCollection(request);
        checkResponseNotFailed(response);
        return (Boolean) response.getData();
    }

    static void createCollection(MilvusServiceClient milvusClient, String collectionName, FieldDefinition fieldDefinition, int dimension) {
        CreateCollectionParam request = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withSchema(CollectionSchemaParam.newBuilder()
                        .addFieldType(FieldType.newBuilder().withName(fieldDefinition.getIdFieldName())
                                .withDataType(DataType.VarChar)
                                .withMaxLength(36).withPrimaryKey(true)
                                .withAutoID(false).build())
                        .addFieldType(FieldType.newBuilder().withName(fieldDefinition.getTextFieldName())
                                .withDataType(DataType.VarChar).withMaxLength(65535).build())
                        .addFieldType(FieldType.newBuilder().withName(fieldDefinition.getMetadataFieldName())
                                .withDataType(DataType.JSON).build())
                        .addFieldType(FieldType.newBuilder().withName(fieldDefinition.getVectorFieldName())
                                .withDataType(DataType.FloatVector).withDimension(dimension).build()).build())
                .build();
        R<RpcStatus> response = milvusClient.createCollection(request);
        checkResponseNotFailed(response);
    }

    static void dropCollection(MilvusServiceClient milvusClient, String collectionName) {
        DropCollectionParam request = CollectionRequestBuilder.buildDropCollectionRequest(collectionName);
        R<RpcStatus> response = milvusClient.dropCollection(request);
        checkResponseNotFailed(response);
    }

    static void createIndex(MilvusServiceClient milvusClient, String collectionName, String vectorFieldName, IndexType indexType, MetricType metricType) {
        CreateIndexParam request = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName(vectorFieldName)
                .withIndexType(indexType)
                .withMetricType(metricType).build();
        R<RpcStatus> response = milvusClient.createIndex(request);
        checkResponseNotFailed(response);
    }

    static void insert(MilvusServiceClient milvusClient, String collectionName, List<InsertParam.Field> fields) {
        InsertParam request = CollectionRequestBuilder
                .buildInsertRequest(collectionName, fields);
        R<MutationResult> response = milvusClient.insert(request);
        checkResponseNotFailed(response);
    }

    static void loadCollectionInMemory(MilvusServiceClient milvusClient, String collectionName) {
        LoadCollectionParam request = CollectionRequestBuilder.buildLoadCollectionInMemoryRequest(collectionName);
        R<RpcStatus> response = milvusClient.loadCollection(request);
        checkResponseNotFailed(response);
    }

    static SearchResultsWrapper search(MilvusServiceClient milvusClient, SearchParam searchRequest) {
        R<SearchResults> response = milvusClient.search(searchRequest);
        checkResponseNotFailed(response);
        return new SearchResultsWrapper(((SearchResults) response.getData()).getResults());
    }

    static QueryResultsWrapper queryForVectors(MilvusServiceClient milvusClient, String collectionName, FieldDefinition fieldDefinition, List<String> rowIds, ConsistencyLevelEnum consistencyLevel) {
        QueryParam request = CollectionRequestBuilder.buildQueryRequest(collectionName, fieldDefinition, rowIds, consistencyLevel);
        R<QueryResults> response = milvusClient.query(request);
        checkResponseNotFailed(response);
        return new QueryResultsWrapper((QueryResults) response.getData());
    }

    static void removeForVector(MilvusServiceClient milvusClient, String collectionName, String expr) {
        R<MutationResult> response = milvusClient.delete(CollectionRequestBuilder.buildDeleteRequest(collectionName, expr));
        checkResponseNotFailed(response);
    }

    private static <T> void checkResponseNotFailed(R<T> response) {
        if (response == null) {
            throw new RequestToMilvusFailedException("Request to Milvus DB failed. Response is null");
        } else if (response.getStatus() != R.Status.Success.getCode()) {
            String message = String.format("Request to Milvus DB failed. Response status:'%d'.%n", response.getStatus());
            throw new RequestToMilvusFailedException(message, response.getException());
        }
    }
}
